#!/bin/bash
# 若想进行编译：./wafer.sh -i input -j json -o output

rm -rf ../result/*
rm -f sha1_host
rm -f *.mlir
rm -rf *.o
# rm -rf *.out
rm -rf *.ll
rm -rf *.s
rm -rf *.cubin
rm -rf *.fatbin
rm -rf _app_cuda_version*
rm -rf _cuobjdump_list_ptx*
rm -rf *.ptx
rm -rf *.ptxas
rm -rf *.txt

rm -rf wafer_gpu.out wafer_run_exe

rm -rf *.json

export WAFER_DIR=/root/wafer_project/wafer-compiler
export WAFER_BUILD_DIR=${WAFER_DIR}/build
export WAFER_FRONTEND=${WAFER_BUILD_DIR}/bin/wafer-frontend
export WAFER_OPT=${WAFER_BUILD_DIR}/bin/wafer-opt
export WAFER_TRANSLATE=${WAFER_BUILD_DIR}/bin/wafer-translate
export LLVM_BIN_DIR=${WAFER_DIR}/3rdparty/llvm/build/bin
export CLANG=${LLVM_BIN_DIR}/clang++
export CICC=/usr/local/cuda/nvvm/bin/cicc

export CXX=/usr/bin/g++
export OPENSSL_DIR=${WAFER_DIR}/3rdparty/openssl/build
export OPENSSL_LIB_DIR=${OPENSSL_DIR}/lib64
export OPENSSL_INC_DIR=${OPENSSL_DIR}/include

INC_DIR=${WAFER_DIR}/include
LLVM_PRJ_DIR=${WAFER_DIR}/3rdparty/llvm
LLVM_CLANG_INC_DIR=${LLVM_PRJ_DIR}/build/lib/clang/17/include
LLVM_LIB_DIR=${LLVM_PRJ_DIR}/build/lib

# 默认参数设置
input_file=""
json_file=""
output_file=""
dump_graph=false
openssl=false
printschedule=OFF

# 编译构建输出目录，需要是相对路径
OUT_DIR=../result/build

rm -rf ../result/*
mkdir -p ${OUT_DIR}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    key="$1"

    case $key in
        -i)
            input_file="$2"
            shift
            shift
            ;;
        -j)
            json_file="$2"
            shift
            shift
            ;;
        -o)
            output_file="$2"
            shift
            shift
            ;;
        -dumpgraph)
            dump_graph=true
            shift
            ;;
        -openssl)
            openssl=true
            shift
            ;;
        -printschedule)
            printschedule=ON
            shift
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

if [ -n "$json_file" ]; then
    # 删除已存在的 design.json 软链接（如果有）
    rm -f ./design.json

    # 创建软链接，名称为 design.json
    ln -s "$json_file" ./design.json
fi

# 根据参数执行相应操作
if $dump_graph; then
    if [ -z "$output_file" ] || [ -z "$input_file" ]; then
        echo "Error: Input file, output file not specified"
        exit 1
    fi
    echo "Dumping graph to $output_file"
    ${WAFER_BUILD_DIR}/bin/wafer-frontend -I ${INC_DIR} -I ${LLVM_CLANG_INC_DIR} -I ${LLVM_PRJ_DIR}/llvm/include -I ${LLVM_PRJ_DIR}/build/include -c "$input_file" --GraphOutputPath "$output_file";
elif $openssl; then
    if [ -z "$output_file" ] || [ -z "$input_file" ]; then
        echo "Error: Input file, output file not specified"
        exit 1
    fi
    ${CXX} -O3 -DWAFER_OPENSSL "$input_file" -I ${OPENSSL_INC_DIR} -L ${OPENSSL_LIB_DIR} -l crypto -L${WAFER_BUILD_DIR}/lib -lwafer_sst_runtime -Wl,-rpath=${WAFER_BUILD_DIR}/lib -o "$output_file"
else
    if [ -z "$input_file" ] || [ -z "$json_file" ] || [ -z "$output_file" ]; then
        echo "Error: Input file, JSON file, or output file not specified"
        exit 1
    fi
    echo "Compiling input file $input_file using JSON file $json_file to output file $output_file"

    ${WAFER_FRONTEND} -I ${INC_DIR} -I ${LLVM_CLANG_INC_DIR} -I ${LLVM_PRJ_DIR}/llvm/include -I ${LLVM_PRJ_DIR}/build/include \
		 -c "$input_file" -o ${OUT_DIR}/crypto_op.mlir
    
    export DESIGN_JSON_FILE=${json_file}
    export PRINT_SCHEDULE=${printschedule}

    ${WAFER_OPT} ${OUT_DIR}/crypto_op.mlir --lower-crypto-to-hyper -o ${OUT_DIR}/crypto_op_hyper.mlir

	${WAFER_OPT} ${OUT_DIR}/crypto_op_hyper.mlir \
		--sst-device-schedule \
		--lower-hyper-sstmemcpy \
		--sst-lower-hyper \
		--lower-crypto-padmessage \
		--lower-cryptosha1-without-SHA1ISA \
		--llvm-request-c-wrappers \
		--convert-scf-to-cf \
		--convert-arith-to-llvm \
		--finalize-memref-to-llvm \
		--convert-vector-to-llvm \
		--lower-sst-to-llvm \
		--convert-func-to-llvm \
		--reconcile-unrealized-casts | \
	${WAFER_TRANSLATE} --wafer-to-llvmir -o ${OUT_DIR}/crypto_op.ll

    # ${WAFER_BUILD_DIR}/bin/wafer-frontend -I ${WAFER_DIR}/3rdparty/llvm/build/lib/clang/17/include  -c "$input_file" --sst --JsonFilePath "$json_file" --wafer-to-llvmir -o ${OUT_DIR}/crypto_op.ll


    ${WAFER_BUILD_DIR}/bin/wafer-frontend -I ${INC_DIR} -I ${LLVM_CLANG_INC_DIR} -I ${LLVM_PRJ_DIR}/llvm/include -I ${LLVM_PRJ_DIR}/build/include --ptx --JsonFilePath "$json_file" -c "$input_file"  -o ${OUT_DIR}/gpulaunch.mlir
    ${WAFER_OPT} ${OUT_DIR}/gpulaunch.mlir -device-schedule -o ${OUT_DIR}/ds.mlir
    if grep -q "gpu.launch" ${OUT_DIR}/ds.mlir; then
        echo "Found 'gpu.launch' in ${OUT_DIR}/ds.mlir."
        ${WAFER_OPT} ${OUT_DIR}/ds.mlir -lower-hyper-memcpy -lower-hyper  -lower-crypto-padmessage -lower-cryptosha1-without-SHA1ISA  -async-to-async-runtime -llvm-request-c-wrappers -async-runtime-ref-counting  -convert-async-to-llvm  -convert-scf-to-cf  -gpu-kernel-outlining  -convert-arith-to-llvm -finalize-memref-to-llvm  -convert-vector-to-llvm -o ${OUT_DIR}/pass.mlir    
        ${WAFER_OPT} ${OUT_DIR}/pass.mlir -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin{chip=sm_30 features=+ptx32}))' > /dev/null
        mv ./kernel.ptx ${OUT_DIR}/kernel.ptx
        gcc -D__CUDA_ARCH__=300 -E -x c++  -DCUDA_DOUBLE_MATH_FUNCTIONS -D__CUDACC__ -D__NVCC__  "-I/usr/local/cuda/bin/../targets/x86_64-linux/include"    -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 -include "cuda_runtime.h" -m64 "./main.cu" > "${OUT_DIR}/main.cpp1.ii" 
        ${CICC} --c++14 --gnu_version=70500 --allow_managed   -arch compute_30 -m64 -ftz=0 -prec_div=1 -prec_sqrt=1 -fmad=1 --include_file_name "main.fatbin.c" -tused -nvvmir-library "/usr/local/cuda/bin/../nvvm/libdevice/libdevice.10.bc" --gen_module_id_file --module_id_file_name "${OUT_DIR}/main.module_id" --orig_src_file_name "./main.cu" --gen_c_file_name "${OUT_DIR}/main.cudafe1.c" --stub_file_name "${OUT_DIR}/main.cudafe1.stub.c" --gen_device_file_name "${OUT_DIR}/main.cudafe1.gpu"  "${OUT_DIR}/main.cpp1.ii" -o "${OUT_DIR}/main.ptx"
        cp ${OUT_DIR}/kernel.ptx ${OUT_DIR}/main.ptx
        fatbinary --create="${OUT_DIR}/main.fatbin" -64 "--image=profile=compute_30,file=${OUT_DIR}/main.ptx" --embedded-fatbin="${OUT_DIR}/main.fatbin.c" 
        gcc -E -x c++ -D__CUDACC__ -D__NVCC__  "-I/usr/local/cuda/bin/../targets/x86_64-linux/include"    -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 -include "cuda_runtime.h" -m64 "./main.cu" > "${OUT_DIR}/main.cpp4.ii" 
        cudafe++ --c++14 --gnu_version=70500 --allow_managed  --m64 --parse_templates --gen_c_file_name "${OUT_DIR}/main.cudafe1.cpp" --stub_file_name "main.cudafe1.stub.c" --module_id_file_name "${OUT_DIR}/main.module_id" "${OUT_DIR}/main.cpp4.ii" 
        gcc -D__CUDA_ARCH__=300 -c -x c++  -DCUDA_DOUBLE_MATH_FUNCTIONS "-I/usr/local/cuda/bin/../targets/x86_64-linux/include"   -m64 -o "${OUT_DIR}/main.o" "${OUT_DIR}/main.cudafe1.cpp" 
        nvlink --arch=compute_30 --register-link-binaries="${OUT_DIR}/dlink.reg.c"  -m64   "-L/usr/local/cuda/bin/../targets/x86_64-linux/lib/stubs" "-L/usr/local/cuda/bin/../targets/x86_64-linux/lib" -cpu-arch=X86_64 "${OUT_DIR}/main.o" 
        fatbinary --create="${OUT_DIR}/dlink.fatbin" -64 -link --embedded-fatbin="${OUT_DIR}/dlink.fatbin.c" 
        gcc -c -x c++ -DFATBINFILE="\"dlink.fatbin.c\"" -DREGISTERLINKBINARYFILE="\"dlink.reg.c\"" -I${OUT_DIR} -D__NV_EXTRA_INITIALIZATION= -D__NV_EXTRA_FINALIZATION= -D__CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__  "-I/usr/local/cuda/bin/../targets/x86_64-linux/include"    -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 -m64 -o "${OUT_DIR}/dlink.o" "/usr/local/cuda/bin/crt/link.stub" 
        g++ -m64 -o ${OUT_DIR}/wafer_gpu.out -Wl,--start-group "${OUT_DIR}/dlink.o" "${OUT_DIR}/main.o"   "-L/usr/local/cuda/bin/../targets/x86_64-linux/lib/stubs" "-L/usr/local/cuda/bin/../targets/x86_64-linux/lib" -lcudadevrt  -L/usr/local/cuda/lib64 -lcudart  -lrt -lpthread  -ldl  -Wl,--end-group
    else
        echo "No 'gpu.launch' found in ds.mlir."
    fi

    ${LLVM_BIN_DIR}/llc -O3 -mtriple=x86_64-unknown-linux-gnu -mattr=+sha ${OUT_DIR}/crypto_op.ll --filetype=obj -o ${OUT_DIR}/host_ctrl.o

    ${CLANG}  -c "$input_file" -I ${INC_DIR} -I ${LLVM_PRJ_DIR}/llvm/include -I ${LLVM_PRJ_DIR}/build/include -o ${OUT_DIR}/example.o

    ${CLANG} ${OUT_DIR}/host_ctrl.o ${OUT_DIR}/example.o \
        -L${WAFER_DIR}/3rdparty/llvm/build/lib -lmlir_c_runner_utils \
        -L${WAFER_BUILD_DIR}/lib -lwafer_sst_runtime \
        -L${WAFER_BUILD_DIR}/lib/CGRAAPI -lwafer_cgra_api \
        -Wl,-rpath=${WAFER_DIR}/3rdparty/llvm/build/lib:${WAFER_BUILD_DIR}/lib \
        -o "$output_file"

    echo "Completed!"
fi

DIR_COPY="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
DIR_RESULT="${DIR_COPY}/../result"
ln -s ${DIR_COPY}/${OUT_DIR}/wafer_gpu.out ${DIR_RESULT}/wafer_gpu.out
ln -s ${DIR_COPY}/wafer_run_exe ${DIR_RESULT}/wafer_run_exe